import torch.cuda
from torch import nn

from common_models.MLP import MLP

device = 'cuda' if torch.cuda.is_available() else 'cpu'


class Expreriment(nn.Module):
    def __init__(self, dataset, args):
        super(Expreriment, self).__init__()
        # 输入维度为图片特征维度，输出维度为属性的数量
        self.attr_classifier = MLP(dataset.feat_dim, 1024, relu=False).to(device)
        self.label_classifier = MLP(1024, args.n_classes, relu=False,
                                    dropout=args.dropout, norm=args.norm).to(device)

    def forward(self, imgs, labels=None, attrs=None):
        # 根据图片特征学习对象的属性
        attr_preds = self.attr_classifier(imgs)  # 属性评分
        label_preds = self.label_classifier(attr_preds)
        softmax = torch.nn.Softmax(dim=1)
        crossEntropyLoss = torch.nn.CrossEntropyLoss()
        loss = crossEntropyLoss(label_preds,labels)
        return loss,attr_preds,label_preds
